import xml.etree.ElementTree as ET
import re
import os
from collections import defaultdict

## 自己定义识别矩形的结构体
class Rec_rectangle:
    def __init__(self, center_x, center_y, x_length, y_length):
        self.center_x = center_x
        self.center_y = center_y
        self.x_length = x_length
        self.y_length = y_length

    def __repr__(self):
        return f"Rec_rectangle(center = ({self.center_x:.2f},{self.center_y:.2f}), x_length = {self.x_length:.2f}, y_length = {self.y_length})"


## 读取svg矢量图把其中直角角度的矩形提取出来，按照左上到右下角进行排序
def parse_svg_rectangles(svg_file):
    tree = ET.parse(svg_file)
    root = tree.getroot()
    ns = {'svg': 'http://www.w3.org/2000/svg'}

    rectangles = []

    for path in root.findall('.//svg:path', ns):
        d = path.attrib.get('d', '')
        # 匹配过程：正则表达式 M x y L x1 y L x2 y2 L x3 y3 Z 或带 Z 的闭合
        paring = list(map(float, re.findall(r'[-+]?[0-9]*\.?[0-9]+', d)))
        if 'Z' not in d and 'z' not in d:
            continue

        if len(paring) == 10:
            points = [(paring[i], paring[i+1]) for i in range(0, 8, 2)]
            xs = [pt[0] for pt in points]
            ys = [pt[1] for pt in points]

            min_x, max_x = min(xs), max(xs)
            min_y, max_y = min(ys), max(ys)

            center_x = (min_x + max_x) / 2
            center_y = (min_y + max_y) / 2
            x_length = abs(max_x - min_x)
            y_length = abs(max_y - min_y)

            ## 筛除掉横纵比太离谱的矩形,这里设置3倍筛选
            length = max(x_length, y_length)
            width = min(x_length, y_length)

            if length / width > 3:
                continue

            if xs.count(xs[0]) + xs.count(xs[2]) == 4 or ys.count(ys[0]) + ys.count(ys[2]) == 4:
                rectangles.append(Rec_rectangle(center_x, center_y, x_length, y_length))

    rectangles.sort(key=lambda r:(r.center_y, r.center_x))
    return rectangles


def group_rects(rects):
    grouped_rects = defaultdict(list)

    for rect in rects:
        x_rounded = round(rect.x_length, 1)
        y_rounded = round(rect.y_length, 1)
        key = (x_rounded, y_rounded)
        grouped_rects[key].append(rect)

    # 输出每组的数量和尺寸
    for size, group in grouped_rects.items():
        print(f"尺寸为 {size} 的矩形有 {len(group)} 个")

    # 存储排序后的分组
    sorted_grouped_rects = {}

    for size, group in grouped_rects.items():
        sorted_group = sorted(group, key=lambda r: (r.center_y, r.center_x))
        sorted_grouped_rects[size] = sorted_group
    return sorted_grouped_rects


def pair_rects_by_direction(rect_group, direction, used):

    pairs = []

    for i, rect1 in enumerate(rect_group):
        if i in used:
            continue

        best_match = None
        min_distance = float('inf')

        for j in range(i + 1, len(rect_group)):
            if j in used:
                continue
            rect2 = rect_group[j]

            dist = float('inf')

            if direction == 'vertical':
                if round(rect1.center_x, 1) == round(rect2.center_x, 1):
                    dist = abs(rect1.center_y - rect2.center_y)

            elif direction == 'horizontal':
                if round(rect1.center_y, 1) == round(rect2.center_y, 1):
                    dist = abs(rect1.center_x - rect2.center_x)

            elif direction is None:
                # 正方形模块处理，双方向匹配，取最近
                ratio = rect1.x_length / rect1.y_length
                if 0.99 <= ratio <= 1.01:
                    dist_h = abs(rect1.center_x - rect2.center_x) if round(rect1.center_y, 2) == round(rect2.center_y, 2) else float('inf')
                    dist_v = abs(rect1.center_y - rect2.center_y) if round(rect1.center_x, 2) == round(rect2.center_x, 2) else float('inf')
                    dist = min(dist_h, dist_v)
                else:
                    continue  # 非正方形时忽略

            if dist < min_distance:
                min_distance = dist
                best_match = j

        if min_distance < 20 and best_match is not None:
            used.add(i)
            used.add(best_match)
            pairs.append((rect1, rect_group[best_match]))

    return pairs


def process_rect_groups(sorted_groups):
    all_pairs = []

    for size, rect_group in sorted_groups.items():
        x_length, y_length = size
        # 复制一份矩形列表，用于匹配后检查未匹配的矩形
        unmatched_rects = rect_group[:]
        used = set()
        if x_length > y_length:
            # Method 1: 长在 x 轴方向，优先纵向配对
            vertical_pairs = pair_rects_by_direction(rect_group, 'vertical', used)
            all_pairs.extend(vertical_pairs)

            # 剩余未使用的再横向配对
            horizontal_pairs = pair_rects_by_direction(rect_group, 'horizontal', used)
            all_pairs.extend(horizontal_pairs)



        elif x_length < y_length:
            # Method 2: 长在 y 轴方向，优先横向配对
            horizontal_pairs = pair_rects_by_direction(rect_group, 'horizontal', used)
            all_pairs.extend(horizontal_pairs)

            # 剩余未使用的再横向配对
            vertical_pairs = pair_rects_by_direction(rect_group, 'vertical', used)
            all_pairs.extend(vertical_pairs)
        else:
            # Method 3: 正方形，横纵都尝试，选最近距离
            pairs = pair_rects_by_direction(rect_group, None, used)
            all_pairs.extend(pairs)


        # 从未匹配的矩形中移除已经成功匹配的矩形
        matched_rects = [rect1 for pair in all_pairs for rect1 in pair]
        unmatched_rects = [rect for rect in unmatched_rects if rect not in matched_rects]

        # 输出未匹配的矩形换方向匹配
        if len(unmatched_rects) > 0:
            print(f'size = {size}尺寸下有{len(unmatched_rects)}个未匹配，分别是：')
            for unmatched in unmatched_rects:
                print(unmatched)
        else:
            print(f"尺寸 {size} 所有矩形都已匹配成功")

    return all_pairs


## 测试生成矩形是否准确  !!!!!  仅供测试使用
def generate_svg(rects, origin_file, output_file = "output_rects.svg"):
    with open(origin_file, 'r', encoding='utf-8') as f:
        content = f.read()

    svg_width = re.search(r'width="([\d.]+)pt"', content)
    svg_height = re.search(r'height="([\d.]+)pt"', content)

    svg = ET.Element('svg',{
        'xmlns': 'http://www.w3.org/2000/svg',
        'width': f'{svg_width}pt',
        'height': f'{svg_height}pt',
        'viewBox': f'0 0 {svg_width} {svg_height}',
        'version': '1.1'
    })

    for rect in rects:
        x = rect.center_x - rect.x_length / 2
        y = rect.center_y - rect.y_length / 2
        ET.SubElement(svg, 'rect', {
            'x': str(x),
            'y': str(y),
            'width': str(rect.x_length),
            'height': str(rect.y_length),
            'fill': 'black',
            'stroke': 'black',
            'stroke-width': '0'
        })
    tree = ET.ElementTree(svg)
    tree.write(output_file, encoding='utf-8', xml_declaration=True)
    print(f"[✓] 已生成 SVG 文件：{output_file}")

## 测试最终配对图案是否正确   !!!! 仅供测试使用
def draw_combined_pairs_svg(pairs, origin_file, output_file="combined_pairs.svg", xml_output_file="combined_pairs_data.xml"):
    with open(origin_file, 'r', encoding='utf-8') as f:
        content = f.read()

    svg_width = re.search(r'width="([\d.]+)pt"', content).group(1)
    svg_height = re.search(r'height="([\d.]+)pt"', content).group(1)

    svg = ET.Element('svg', {
        'xmlns': 'http://www.w3.org/2000/svg',
        'width': f'{svg_width}pt',
        'height': f'{svg_height}pt',
        'viewBox': f'0 0 {svg_width} {svg_height}',
        'version': '1.1'
    })

    # ====== 新增：创建 XML根节点 ======
    root_xml = ET.Element('CombinedPairs')

    for rect1, rect2 in pairs:
        # 分别计算两个矩形的左上角和右下角坐标
        x1_min = rect1.center_x - rect1.x_length / 2
        y1_min = rect1.center_y - rect1.y_length / 2
        x1_max = rect1.center_x + rect1.x_length / 2
        y1_max = rect1.center_y + rect1.y_length / 2

        x2_min = rect2.center_x - rect2.x_length / 2
        y2_min = rect2.center_y - rect2.y_length / 2
        x2_max = rect2.center_x + rect2.x_length / 2
        y2_max = rect2.center_y + rect2.y_length / 2

        # 合并区域的左上角和右下角
        left = min(x1_min, x2_min)
        top = min(y1_min, y2_min)
        right = max(x1_max, x2_max)
        bottom = max(y1_max, y2_max)

        width = right - left
        height = bottom - top
        center_x = (left + right) / 2
        center_y = (top + bottom) / 2

        # 添加合并的大矩形（黑色）
        ET.SubElement(svg, 'rect', {
            'x': str(left),
            'y': str(top),
            'width': str(width),
            'height': str(height),
            'fill': 'black',
            'stroke': 'black',
            'stroke-width': '0'
        })

        # 添加红色中心点（小圆点）
        ET.SubElement(svg, 'circle', {
            'cx': str(center_x),
            'cy': str(center_y),
            'r': '1.0',
            'fill': 'red'
        })

        # === 将数据写入XML节点中 ===
        pair_elem = ET.SubElement(root_xml, 'Pair')
        ET.SubElement(pair_elem, 'CenterX').text = str(center_x)
        ET.SubElement(pair_elem, 'CenterY').text = str(center_y)
        ET.SubElement(pair_elem, 'Width').text = str(width)
        ET.SubElement(pair_elem, 'Height').text = str(height)

    tree = ET.ElementTree(svg)
    tree.write(output_file, encoding='utf-8', xml_declaration=True)
    print(f"[✓] 已生成合并矩形 SVG 文件：{output_file}")

    tree_xml = ET.ElementTree(root_xml)
    tree_xml = ET.ElementTree(root_xml)
    tree_xml.write(xml_output_file, encoding='utf-8', xml_declaration=True)
    print(f"[✓] 已生成中心点与长宽数据 XML 文件：{xml_output_file}")

if __name__ == "__main__":
    svg_file = r'C:\pycharm_proj\PCB_Detection\svg\32.svg'
    rects = parse_svg_rectangles(svg_file)
    sorted_groups = group_rects(rects)
    matched_pairs = process_rect_groups(sorted_groups)

    generate_svg(rects, r'C:\pycharm_proj\PCB_Detection\svg\32.svg')
    draw_combined_pairs_svg(matched_pairs, r'C:\pycharm_proj\PCB_Detection\svg\32.svg')